from typing import List, Any, Dict, Tuple
import re
import json
import ast
from langchain.agents.structured_chat.output_parser import StructuredChatOutputParser
from langchain.schema import Generation
from langchain_core.output_parsers.base import T

json_pattern = re.compile(r"```json\s+(\W.*?)```", re.DOTALL)


class BaseOutputParser(StructuredChatOutputParser):
    def parse_result(self, result: List[Generation], *, partial: bool = False) -> Tuple[T, Dict]:
        parsed_content = self.parse(result[0])
        usage = get_usage_from_generation(result[0])
        return parsed_content, usage

    def parse(self, generation: Generation):
        return super().parse(generation.text)


class JSONListOutputParser(BaseOutputParser):
    def parse(self, generation: Generation):
        text = generation.text
        try:
            json_obj = json.loads(text)
            return json_obj
        except Exception as e:
            pass

        try:
            action_match = json_pattern.search(text)
            if action_match is not None:
                try:
                    action = json.loads(action_match.group(1).strip(), strict=False)
                except Exception as e:
                    try:
                        action = ast.literal_eval(action_match.group(1).strip())
                    except Exception as e1:
                        raise e1
                return action if isinstance(action, list) else [action]
            return None
        except Exception as e:
            raise e


class JSONObjOutputParser(BaseOutputParser):
    def parse(self, generation: Generation):
        text = generation.text
        try:
            json_obj = json.loads(text)
            return json_obj
        except Exception as e:
            pass

        try:
            action_match = json_pattern.search(text)
            if action_match is not None:
                try:
                    action = json.loads(action_match.group(1).strip(), strict=False)
                except Exception as e:
                    try:
                        action = ast.literal_eval(action_match.group(1).strip())
                    except Exception as e1:
                        print(f"Error: {e}")
                        raise e1
                return action
            return None
        except Exception as e:
            raise e


def get_usage_from_generation(generation):
    """
    从generation的AImessage中获取usage信息

    Args:
        generation: 包含AI消息的生成对象

    Returns:
        dict: 包含usage信息的字典，如果没有找到则返回None
    """
    try:
        if hasattr(generation, 'message') and generation.message:
            if hasattr(generation.message, 'usage_metadata') and generation.message.usage_metadata:
                return generation.message.usage_metadata

        # 如果直接在generation上有usage_metadata
        if hasattr(generation, 'usage_metadata') and generation.usage_metadata:
            return generation.usage_metadata

        return {}
    except Exception as e:
        print(f"获取usage信息时出错: {e}")
        return {}
